# coding=utf-8
import math
import tushare as ts
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import talib
import pandas as pd
from datetime import datetime, date
from sklearn import datasets
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression

matplotlib.rcParams['axes.unicode_minus']=False
plt.rcParams['font.sans-serif']=['SimHei']
ts.set_token('6557f39ddf25c764e5124721450abfd1981f5ff98274f68caf9acd9f')
pro = ts.pro_api()


############################读取数据类###################################
class readData:
    def read_index_daily(self,code,star,end):#指数数据
        dsb = pro.index_daily(ts_code=code, start_date=star, end_date=end,fields='ts_code,trade_date,close,change')#默认读取三个数据
        return dsb
    def read_daily(self,code,star,end):
        dsc1 = pro.daily(ts_code=code, start_date=star, end_date=end,fields='ts_code,trade_date,close')
        return dsc1
    def read_CPI(self,star,end):#时间格式start_cpi='201609'
        dc=pro.cn_cpi(start_m=star, end_m=end,fields='month,nt_yoy')
        return dc
    def read_GDP(self,star,end):#时间格式star='2016Q4'
        df1 = pro.cn_gdp(start_q=star, end_q=end,fields='quarter,gdp_yoy')
        return df1
    def read_bond(self,code,star,end):
        df=pro.cb_daily(ts_code=code,start_date=star,end_date=end)
    def read_base(self,code):
        df=pro.query('daily_basic', ts_code=code,fields='close,ts_code,pb,total_mv,trade_date')
        return df
#######################################################################################

start_time='20180226'#发布GDP需要时间，我们延迟1个月，即第一季度的GDP4月份才发布。
end_time="20180521"
dc=readData()
dsc1=readData()
dsb1=readData()
###################计算alpha的值###########################################################
def alpha_fun(code):
    dsb=dsb1.read_base(code) .fillna(0)
    dsc=dsc1.read_index_daily('000300.SH',start_time,end_time)
    dsc.set_index(['trade_date'],inplace=True)
    dsb.set_index(['trade_date'],inplace=True)
    df=pd.merge(dsc, dsb, on='trade_date').fillna(0)
    R=np.reshape(   np.array([df.close_y]) , (-1,1) )
    R_f=np.reshape(     np.array([  (df.change/(df.close_x.shift(-1))).fillna(0)  ])   , (-1,1)      )#用0 填充nan
    HMI=np.reshape(    np.array([ (1/df.pb).fillna(0)    ])                              , (-1,1) )
    SMB=np.reshape(    np.array([ df.total_mv]) , (-1,1) )
    X=np.concatenate((    R_f-4/252, HMI,SMB   ),axis=1)
    y1=np.reshape(R,(1,-1)).T
    X_train, X_test, y_train, y_test = train_test_split(X, y1, test_size=0.3, random_state=0)
    linear = LinearRegression()
    linear.fit(X_train, y_train)
    alpha=linear.intercept_-4/252
    return alpha,linear.intercept_ ,linear.coef_,linear.score(X_test, y_test),df

############################计算篮子里的每只股票净值###############################################
def jinzhi_fun(code,start_time,end_time):#计算净值
    df=pro.query('daily', ts_code=str(code), start_date=start_time, end_date=end_time,fields='')
    df.index=pd.to_datetime(df.trade_date,format="%Y-%m-%d")
    jinzhi=(df.close/df.close[-1])
    return jinzhi

######################选股条件####################################################
co=pro.query('daily_basic', ts_code="",trade_date="20180226",fields='ts_code')
code_list=[]
N=30#股票池
k=0
JZ=0#组合净值
cum=0
for i in co.ts_code.values[0:N]:
    try:
        if alpha_fun(i)[0]<0:
            k+=1
            JZ=jinzhi_fun( str(i) ,start_time,end_time)    +JZ
            JZ=JZ.fillna( method='bfill')#用前后数据特征填充
    except ValueError:
        pass
    continue

#########################计算权重收益净值#########################################
JZ=JZ.sort_index(axis=0,ascending=True)
JZ_avg=JZ/k#平均净值
cum=np.cumprod(JZ_avg)# 累计收益

####################计算收益率函数，如沪深300#####################################
def JZ_function(code,star,end):
    df12 = pro.index_daily( ts_code=code, start_date=star, end_date=end)
    df12=df12.sort_index( )
    df12.index=pd.to_datetime(df12.trade_date,format='%Y-%m-%d')#设置日期索引
    ret12=df12.close/df12.close[-1]
    return ret12.dropna()

#############################策略的年化统计######################################
def Tongji(jz):
    NH=(jz.dropna()[-1]-1)*100*252/len(jz.index)
    BD=np.std(jz)*100*np.sqrt(252)
    SR=(NH-400/252)/BD
    return_list=jz.dropna()
    MHC=((np.maximum.accumulate(return_list) - return_list) / np.maximum.accumulate(return_list)).max()*100
    print("年化收益率：{:.2f}%:，年化夏普率：{:.2f},波动率为：{:.2f}%,最大回撤：{:.2f}%".format( NH,SR,BD,MHC))


############################################################################
if __name__=="__main__":
    cum12=JZ_function('000300.SH',start_time,end_time)
    Tongji(JZ_avg)
    plt.plot(cum12,label="沪深300",color='b')
    plt.plot(JZ_avg,label="股票组合",color='r')
    plt.title("alpha股+指期对冲策略")
    plt.legend()
